import os

# from langchain.chains import LLMChain  # 使用 LCEL 而不用显式导入 LLMChain
from langchain_openai import ChatOpenAI  # for OpenAI, ZhipuAI
from langchain_community.llms import Tongyi
from langchain_core.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

from utils import LOG, generate_token

class TranslatorChain:
    def __init__(self, model_type: str = "OpenAI", model: str = "gpt-3.5-turbo"):
        # 构造系统指令Prompt
        system_template = """
        你是一个翻译专家，你的任务是对用户输入的文本进行翻译，请将其翻译成 {target_language}，对于用户给定的输入只需要进行翻译，不要做其他任务。

        首先先判断用户输入的内容是什么类型，是 text 还是 table，请注意：用户的输入只有可能是这两种类型
            - 如果用户输入的内容类型是 text，则只参考示例1
            - 如果用户输入的内容类型是 table，则只参考示例2

        示例1：
        target_language = 中文
        input: This is a sample text for translation.
        expected_out: 这是一个用于翻译的示例文本。
        
        示例2：
        target_language = 中文
        input: 
            Fruit   Color   Price (USD)
            Apple   Red     1.20
            Banana  Yellow  0.50
            Orange  Orange  0.80
        output:
            | 水果 | 颜色 | 价格（美元） |
            | ---- | ---- | ------------ |
            | 苹果 | 红色 | 1.20         |
            | 香蕉 | 黄色 | 0.50         |
            | 橙子 | 橙色 | 0.80         |
        
        think it step by step.
        只返回用户输入内容翻译后的文本，不要返回其他多余的内容！
        只返回用户输入内容翻译后的文本，不要返回其他多余的内容！
        """
        system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)

        # 构造待翻译文本Prompt
        human_template = "{text}"
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

        # 构造 ChatPrompt
        chat_prompt_template = ChatPromptTemplate.from_messages(
            [system_message_prompt, human_message_prompt]
        )

        # 根据不同的 model_type 构造不同的 TranslatorChain
        if model_type == "OpenAI":
            chat_model = ChatOpenAI(model_name=model)
        elif model_type == "ZhipuAI":
            chat_model = ChatOpenAI(
                model_name=model,
                openai_api_base='https://open.bigmodel.cn/api/paas/v4',
                openai_api_key=generate_token(os.getenv('ZHIPUAI_API_KEY'), 24*60*60),  # 设置 token 的过期时间为 1 天
            )
        elif model_type == "Tongyi":
            chat_model = Tongyi(model_name=model)

        self.chain = chat_prompt_template | chat_model  # LCEL

    def invoke(self, text: str, target_language: str) -> (str, bool):
        try:
            result = self.chain.invoke({
                "text": text,
                "target_language": target_language,
            })
        except Exception as e:
            LOG.error(f"An error occurred during translation: {e}")
            return "", False
        
        return result, True